"""
广义优势估计（GAE）计算

实现 Generalized Advantage Estimation (Schulman et al., 2016)
"""

from __future__ import annotations

import torch


def compute_gae(
    rewards: torch.Tensor,
    values: torch.Tensor,
    dones: torch.Tensor,
    gamma: float = 0.99,
    gae_lambda: float = 0.95,
    next_value: float = 0.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    计算广义优势估计（GAE）
    
    GAE 通过平衡偏差和方差来估计优势函数，使用 lambda 参数来控制
    TD(lambda) 估计器的权重。
    
    Args:
        rewards: 奖励序列 [T]
        values: 价值估计序列 [T]
        dones: 终止标志序列 [T] (1.0表示终止，0.0表示未终止)
        gamma: 折扣因子
        gae_lambda: GAE lambda参数（0-1之间，0为TD(0)，1为Monte Carlo）
        next_value: 最后一个状态之后的价值估计（用于bootstrap）
        
    Returns:
        advantages: 优势估计 [T]
        returns: 目标回报 [T]
        
    Reference:
        Schulman et al., "High-Dimensional Continuous Control Using 
        Generalized Advantage Estimation", 2016
        https://arxiv.org/abs/1506.02438
    """
    advantages = []
    gae = 0.0
    
    # 从后向前计算GAE
    # δ_t = r_t + γ * V(s_{t+1}) * (1 - done_t) - V(s_t)
    # A_t = δ_t + γ * λ * (1 - done_t) * A_{t+1}
    for t in reversed(range(len(rewards))):
        if t == len(rewards) - 1:
            next_value_t = next_value
        else:
            next_value_t = values[t + 1]
        
        # TD error
        delta = rewards[t] + gamma * next_value_t * (1 - dones[t]) - values[t]
        
        # GAE advantage
        gae = delta + gamma * gae_lambda * (1 - dones[t]) * gae
        advantages.insert(0, gae)
    
    advantages = torch.FloatTensor(advantages)
    returns = advantages + values
    
    return advantages, returns


def normalize_advantages(advantages: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    """
    标准化优势估计
    
    标准化可以稳定训练并改善收敛性。
    
    Args:
        advantages: 优势估计 [batch_size]
        eps: 数值稳定性参数
        
    Returns:
        normalized_advantages: 标准化后的优势 [batch_size]
    """
    return (advantages - advantages.mean()) / (advantages.std() + eps)
